# 对AGA (RISAdaptiveGeneticAlgorithm) 的消融实验(AblationStudy) 只有 Mutation

import random
import numpy as np

from beam.util_psll import get_psll_by_phase


class AblationAGAMutation():
    __bit_num = 0  # 比特数
    __beam_num = 0  # 波束数

    def __init__(self, bit_num, beam_num, population_size=50, num_generations=100, num_parents=10, mutation_rate=0.1):
        # 初始化阵列相关
        self.__bit_num = bit_num
        self.__beam_num = beam_num
        # 初始化遗传算法相关
        self.population_size = population_size
        self.num_generations = num_generations
        self.num_parents = num_parents
        self.mutation_rate = mutation_rate
        self.population = None
        self.best_individual = None
        self.best_fitness = None
        self.best_fitness_history = []  # 保存每一代的最佳适应度值
        self.best_individual_history = []  # 保存每一代的最佳适应度值的个体
        # 新增自适应参数
        self.adaptive_params = {
            'pc_min': 0.6,  # 最小交叉概率[7](@ref)
            'pc_max': 0.9,  # 最大交叉概率[7](@ref)
            'pm_base': 0.05,  # 基础变异率[5](@ref)
            'pm_max': 0.2,  # 最大变异率[5](@ref)
            'diversity_threshold': 0.1  # 多样性阈值[6](@ref)
        }

    def initialize_population(self, phase_mix_init):
        """初始化种群"""
        self.population = [phase_mix_init] * self.population_size

    def fitness(self, phase):
        return get_psll_by_phase(phase, self.__bit_num, self.__beam_num)

    def selection(self):
        """选择操作"""
        fitness_scores = [self.fitness(individual) for individual in self.population]
        sorted_indices = np.argsort(fitness_scores)  # 从低到高排序
        selected_parents = [self.population[i] for i in sorted_indices[:self.num_parents]]
        # 获取选中的父代个体对应的适应度值
        selected_fitness_scores = [fitness_scores[i] for i in sorted_indices[:self.num_parents]]
        return selected_parents, selected_fitness_scores

    def crossover(self, parents, offspring_size):
        """交叉操作"""
        offspring = []
        for _ in range(offspring_size):
            parent1, parent2 = random.sample(parents, 2)
            parent1 = np.array(parent1)
            parent2 = np.array(parent2)
            # 获取数组的形状
            shape = parent1.shape
            # 将 parent1 和 parent2 扁平化为一维数组
            flat_parent1 = parent1.flatten()
            flat_parent2 = parent2.flatten()
            # 随机选择一个切割位置
            cut_point = random.randint(0, len(flat_parent1))
            # 生成子代
            child = np.concatenate((flat_parent1[:cut_point], flat_parent2[cut_point:]))
            # 将子代重新塑形为原始的二维数组形状
            child = child.reshape(shape)
            child = child.tolist()
            #
            offspring.append(np.array(child))
        return offspring

    def adaptive_mutation(self, offspring):
        """自适应高斯变异"""
        fitness = np.array([self.fitness(ind) for ind in offspring])
        mean_fitness = np.mean(fitness)

        for i in range(len(offspring)):
            # 基于适应度的变异率调整
            if fitness[i] < mean_fitness:
                current_pm = self.adaptive_params['pm_base'] + \
                             (self.adaptive_params['pm_max'] - self.adaptive_params['pm_base']) * \
                             (1 - fitness[i] / mean_fitness)
            else:
                current_pm = self.adaptive_params['pm_base']

            # 应用自适应高斯变异
            mask = np.random.rand(*offspring[i].shape) < current_pm
            noise = np.random.normal(0, 3, offspring[i].shape)  # 均值0，标准差3
            offspring[i] = np.clip(offspring[i] + mask * noise, -180, 180)
        return offspring

    def run(self, phase_mix_init, logger):
        logger.info("[AblationAGAMutation] population_size=%d, num_generations=%d, num_parents=%d, mutation_rate=%d"
                    % (self.population_size, self.num_generations, self.num_parents, self.mutation_rate))
        # """初始化返回值"""
        self.best_fitness = self.fitness(phase_mix_init)
        self.best_individual = phase_mix_init
        self.best_fitness_history = []
        self.best_individual_history = []
        # """运行遗传算法 -- 初始化阶段"""
        self.initialize_population(phase_mix_init)
        # """运行遗传算法 -- 搜索阶段"""
        for generation in range(self.num_generations):
            # 选择操作
            selected_parents, selected_fitness_scores = self.selection()
            # 交换操作
            offspring = self.crossover(selected_parents, self.population_size - self.num_parents)
            # 消融点: 自适应变异
            offspring = self.adaptive_mutation(offspring)
            # 计算后代的适应度值
            offspring_fitness_scores = [self.fitness(individual) for individual in offspring]
            # 更新种群
            self.population = selected_parents + offspring
            # 合并父代和后代的适应度值
            all_fitness_scores = selected_fitness_scores + offspring_fitness_scores
            #
            # 找到当前代的最佳个体
            this_best_index = np.argmin(all_fitness_scores)  # 找到最佳适应度值的索引
            if this_best_index < len(selected_parents):
                this_best_individual = selected_parents[this_best_index]  # 最佳个体来自父代
            else:
                this_best_individual = offspring[this_best_index - len(selected_parents)]  # 最佳个体来自后代
            this_best_fitness = all_fitness_scores[this_best_index]  # 最佳适应度值
            #
            # 更新最佳适应度
            if this_best_fitness < self.best_fitness:
                self.best_fitness = this_best_fitness
                self.best_individual = this_best_individual
            # 记录最佳适应度曲线
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)
            #
            # logger.info("generation=%d: self.best_fitness=%f, self.best_individual:%s"
            #             % (generation, self.best_fitness, self.best_individual))
            logger.info("generation=%d: self.best_fitness=%f" % (generation, self.best_fitness))
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history

